//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift open source project
//
// Copyright (c) 2025 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See http://swift.org/LICENSE.txt for license information
// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

import SWBCore
import SWBUtil
import SWBMacro
import Foundation
import SWBProtocol

final class HeadermapVFSTaskProducer: StandardTaskProducer, TaskProducer {

    private let targetContexts: [TaskProducerContext]

    init(context globalContext: TaskProducerContext, targetContexts: [TaskProducerContext]) {
        self.targetContexts = targetContexts
        super.init(globalContext)
    }

    func generateTasks() async -> [any PlannedTask] {
        let components = context.globalProductPlan.planRequest.buildRequest.parameters.action.buildComponents
        guard components.contains("build") || components.contains("exportLoc") || (components.contains("api")) else {
            return []
        }

        var tasks = [any PlannedTask]()
        do {
            let vfsContentsByPath = try Dictionary(try await targetContexts.concurrentMap(maximumParallelism: 100) { (targetContext: TaskProducerContext) async throws -> (Path, ByteString)? in
                let targetScope = targetContext.settings.globalScope
                let vfsSetting = targetScope.evaluate(BuiltinMacros.CPP_HEADERMAP_PRODUCT_HEADERS_VFS_FILE)
                guard targetContext.configuredTarget?.target is SWBCore.StandardTarget, targetScope.evaluate(BuiltinMacros.USE_HEADERMAP), !vfsSetting.isEmpty else {
                    return nil
                }
                let vfsPath = self.context.makeAbsolute(vfsSetting)

                let contents = try await targetContext.constructVFSContents()
                return (vfsPath, contents)
            }.compactMap { $0 }, uniquingKeysWith: { first, second in
                guard first == second else {
                    throw StubError.error("Unexpected difference in VFS content.\nFirst: \(first.asString)\nSecond:\(second.asString)")
                }
                return first
            })

            for (vfsPath, contents) in vfsContentsByPath {
                await appendGeneratedTasks(&tasks) { delegate in
                    let orderingInputs = delegate.buildDirectories.sorted().filter { $0.isAncestor(of: vfsPath) }.map { delegate.createBuildDirectoryNode(absolutePath: $0) }
                    context.writeFileSpec.constructFileTasks(CommandBuildContext(producer: context, scope: context.settings.globalScope, inputs: [], output: vfsPath, commandOrderingInputs: orderingInputs), delegate, contents: contents, permissions: nil, preparesForIndexing: true, additionalTaskOrderingOptions: [.immediate])
                }
            }

            await appendGeneratedTasks(&tasks) { delegate in
                delegate.createGateTask(inputs: vfsContentsByPath.keys.sorted().map { delegate.createNode($0) }, output: delegate.createVirtualNode("WorkspaceHeaderMapVFSFilesWritten"), name: "WorkspaceHeaderMapVFSFilesWritten")
            }
        } catch {
            self.context.error(error.localizedDescription)
        }

        return tasks
    }
}

extension TaskProducerContext {
    func constructVFSContents() async throws -> ByteString {
        let vfs = VFS()
        // For VFS purposes, always map to the temporary copies of the module maps, they may have been autogenerated or modified from the sources.
        // We only apply over targets participating in the build, to match Xcode.
        await forEachBuildProductToSourceMapping(mapToTempFilesOnly: true) { target, isActiveTarget, builtFile, originalSourceFile in
            if isActiveTarget {
                assert(builtFile.isAbsolute)
                assert(originalSourceFile.isAbsolute)
                vfs.addMapping(builtFile, externalContents: originalSourceFile)
            }
        }
        return try PropertyListItem(vfs.toVFSOverlay()).asJSONFragment()
    }

    func forEachBuildProductToSourceMapping(mapToTempFilesOnly: Bool, _ body: ((target: SWBCore.Target, isActiveTarget: Bool, builtFile: Path, originalSource: Path)) throws -> Void) async rethrows {

        let index = await workspaceContext.headerIndex

        // Get the project header info.
        //
        // FIXME: These conditions should be impossible, but see: <rdar://problem/31248574> SWBBuildService crashed in HeadermapTaskProducer.constructVFSContents()
        let project = workspaceContext.workspace.project(for: configuredTarget!.target)
        guard let projectInfo = index.projectHeaderInfo[project] else {
            fatalError("unexpected missing header info for \(project.name)")
        }

        let activeTargets = Set(globalProductPlan.allTargets.map { $0.target })

        let hasEnabledIndexBuildArena = settings.globalScope.evaluate(BuiltinMacros.INDEX_ENABLE_BUILD_ARENA)
        let effectivePlatformName = settings.globalScope.evaluate(BuiltinMacros.EFFECTIVE_PLATFORM_NAME)
        var settingsByTarget: [SWBCore.Target: Settings] = [:]
        for configuredTarget in globalProductPlan.allTargets {
            guard workspaceContext.workspace.project(for: configuredTarget.target) == project else { continue }
            let settings = globalProductPlan.getTargetSettings(configuredTarget)
            guard settings.globalScope.evaluate(BuiltinMacros.EFFECTIVE_PLATFORM_NAME) == effectivePlatformName else { continue }
            settingsByTarget[configuredTarget.target] = settings
        }

        for (target, targetInfo) in projectInfo.targetHeaderInfo.sorted(byKey: \.name) {
            // Index build needs the VFS to have consistent contents independent of the list of targets that are built.
            let isActiveTarget = hasEnabledIndexBuildArena || activeTargets.contains(target)

            let settings: Settings
            if let targetSettings = settingsByTarget[target] {
                settings = targetSettings
            } else {
                // Either the configured target platform doesn't match, or we're in an index build and somehow missing a target

                let unconfiguredSettings = globalProductPlan.getUnconfiguredTargetSettings(target, viewedFrom: configuredTarget!)

                if hasEnabledIndexBuildArena {
                    // If the only supported platform is unknown (ie. not in the platform registry), then we end up getting back settings for the default platform. This could match the effective platform, but we don't want to include mappings for it - we'll never run any tasks associated with it, so any mappings pointing into the build directory (eg. the module map) will be missing.
                    // Usually this isn't a problem, but if multiple targets are defined with the same product name then the generated module map can be overwritten and thus point to these missing files.
                    // To avoid this case, check that the platform we get back is actually in the supported platforms.
                    // TODO: This could all be avoided if the mappings were target specific and built from the closure of dependent targets (rdar://72912847)
                    // FIXME: The comments above might be out of date now that there is no longer a default platform.
                    let supportedPlatforms = unconfiguredSettings.globalScope.evaluate(BuiltinMacros.SUPPORTED_PLATFORMS)
                    if supportedPlatforms.count > 0 {
                        guard let platformName = unconfiguredSettings.platform?.name else {
                            continue
                        }
                        if !supportedPlatforms.contains(platformName) {
                            continue
                        }
                    }

                    if unconfiguredSettings.globalScope.evaluate(BuiltinMacros.EFFECTIVE_PLATFORM_NAME) != effectivePlatformName {
                        continue
                    }
                }

                settings = unconfiguredSettings
            }

            let scope = settings.globalScope
            let headerDestPaths = TargetHeaderInfo.builtProductDestDirs(scope: scope, workingDirectory: defaultWorkingDirectory)
            let builtProductsDir = headerDestPaths.basePath
            let publicHeadersFolderPseudoPath = headerDestPaths.publicPath
            let privateHeadersFolderPseudoPath = headerDestPaths.privatePath

            // Add the public and private headers.
            func addEntry(_ fileRef: SWBCore.FileReference, installDir: Path) throws {
                // Compute the header path.
                //
                // FIXME: This isn't the correct file resolver to use.
                let path = settings.filePathResolver.resolveAbsolutePath(fileRef)

                // Compute the installed header path.
                let installPath = installDir.join(path.basename)

                try body((target, isActiveTarget, installPath, path))
            }
            for entry in targetInfo.publicHeaders {
                try addEntry(entry.fileReference, installDir: publicHeadersFolderPseudoPath)
            }
            for entry in targetInfo.privateHeaders {
                try addEntry(entry.fileReference, installDir: privateHeadersFolderPseudoPath)
            }

            // If this product defines a module, add some additional mappings.
            //
            // FIXME: This does not belong here, this is specific to the module map creation and the Swift compiler, so they shouldn't need to be here inside common code. However, we don't currently expose any API for unrelated tasks to contribute to the VFS. We should consider providing one so that this can migrate elsewhere.
            if scope.evaluate(BuiltinMacros.DEFINES_MODULE) {
                let contentsFolderString = scope.evaluate(BuiltinMacros.CONTENTS_FOLDER_PATH).str
                let wrapperNameString = scope.evaluate(BuiltinMacros.WRAPPER_NAME).str

                let moduleInfo = GlobalProductPlan.computeModuleInfo(workspaceContext: globalProductPlan.planRequest.workspaceContext, target: target, settings: settings, diagnosticHandler: { message, location, component, essential in
                    if essential {
                        error(message, location: location, component: component)
                    }
                })
                if let moduleInfo = moduleInfo {
                    let moduleMapSourcePath = mapToTempFilesOnly ? moduleInfo.moduleMapPaths.tmpPath : moduleInfo.moduleMapPaths.sourcePath
                    if !moduleMapSourcePath.isEmpty {
                        let builtFile = Path(moduleInfo.moduleMapPaths.builtPath.str.replacingOccurrences(of: contentsFolderString, with: wrapperNameString))
                        let sourceFile = Path(moduleMapSourcePath.str.replacingOccurrences(of: contentsFolderString, with: wrapperNameString))
                        try body((target, isActiveTarget, builtFile, sourceFile))
                    }

                    if let privateModuleMapPaths = moduleInfo.privateModuleMapPaths {
                        let privateModuleMapSourcePath = mapToTempFilesOnly ? privateModuleMapPaths.tmpPath : privateModuleMapPaths.sourcePath
                        if !privateModuleMapSourcePath.isEmpty {
                            let builtFile = Path(privateModuleMapPaths.builtPath.str.replacingOccurrences(of: contentsFolderString, with: wrapperNameString))
                            let sourceFile = Path(privateModuleMapSourcePath.str.replacingOccurrences(of: contentsFolderString, with: wrapperNameString))
                            try body((target, isActiveTarget, builtFile, sourceFile))
                        }
                    }

                    // Only invoke the closure for this file if the file is really going to be generated.
                    if moduleInfo.exportsSwiftObjCAPI {
                        let generatedSwiftHeaderPath = SwiftCompilerSpec.generatedObjectiveCHeaderOutputPath(scope)
                        let swiftHeaderName = scope.evaluate(BuiltinMacros.SWIFT_OBJC_INTERFACE_HEADER_NAME) // if exportsSwiftObjCAPI, then swiftHeaderName is not empty
                        // FIXME: The calculation of swiftHeaderPath below seems like an unfortunate hard-coding of an algorithm that should be elsewhere; it comes from commit ce4151bf9e in HeadermapTaskProducer.swift originally.
                        // FIXME: It's also not clear under what conditions this calculation will be different from SwiftCompilerSpec.generatedObjectiveCHeaderOutputPath(), but if it is (now or someday) we invoke the closure.
                        let swiftHeaderPath = builtProductsDir.join(publicHeadersFolderPseudoPath).join(swiftHeaderName)
                        try body((target, isActiveTarget, generatedSwiftHeaderPath, swiftHeaderPath))
                    }
                }
            }
        }
    }
}
